import os

from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset
from torch.utils.data import SubsetRandomSampler
from torchvision.datasets import ImageFolder
from preprocess.transforms import ToTensor,Compose,Center_Crop
from batchgenerators.transforms.abstract_transforms import Compose
from torchvision import transforms
import SimpleITK as sitk
class TestDataset(Dataset):
    def __init__(self, image_path_list):
        # generate image path list

        self.image_path_list = image_path_list

        self.transforms = Compose([
            # transforms.CenterCrop((128, 128)),
            ToTensor(),


        ])

    def __getitem__(self, index):
        # print (image_path,mask_path)
        image_path,mask_path = self.image_path_list[index]

        
        image = sitk.ReadImage(image_path)
        mask = sitk.ReadImage(mask_path)

        # get label from dirname
        label = int(image_path.split("/")[1])
        label = torch.tensor(label)


        input = torch.stack([image,mask],0)
        input = input /255.0

        
        input = self.transforms(input)
        input = transforms.ToTensor(input)
        return input,label


    def __len__(self):
        return len(self.image_path_list)
